import copy
import time
from random import  randrange,sample
import networkx as nx
import numpy as np
import util

SLOT_NUM = 50
K = 5
TASK_SLOT = 4
TASK_NUM = 120

def onehot_task(task):
    res=np.zeros((21,14))
    res[:,task[0]]=1
    res[:,task[1]]=1
    return res


def create_nsfnet_graph():
    Gbase = nx.Graph()
    Gbase.add_nodes_from([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13])
    edges=[(0, 1), (0, 2), (0, 3), (1, 2), (1, 7), (2, 5), (3, 8), (3, 4), (4, 5), (4, 6), (5, 12), (5, 13),
         (6, 7), (7, 10), (8, 9), (8, 11), (9, 10), (9, 12), (10, 11), (10, 13), (11, 12)]
    Gbase.add_edges_from(edges)
    i = 0
    for u, v, d in Gbase.edges(data=True):
        d['slots'] = np.zeros(SLOT_NUM)
        d['code'] = i
        i += 1
    return Gbase
def create_geant2_graph():
    Gbase = nx.Graph()
    Gbase.add_nodes_from([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23])
    Gbase.add_edges_from(
        [(0, 1), (0, 2), (1, 3), (1, 6), (1, 9), (2, 3), (2, 4), (3, 6), (4, 7), (5, 3),
         (5, 8), (6, 9), (6, 8), (7, 11), (7, 8), (8, 11), (8, 20), (8, 17), (8, 18), (8, 12),
         (9, 10), (9, 13), (9, 12), (10, 13), (11, 20), (11, 14), (12, 13), (12,19), (12,21),
         (14, 15), (15, 16), (16, 17), (17,18), (18,21), (19, 23), (21,22), (22, 23)])
    i = 0
    for u, v, d in Gbase.edges(data=True):
        d['slots'] = np.zeros(SLOT_NUM)
        d['code'] = i
        i += 1
    return Gbase

def create_small_top():
    Gbase = nx.Graph()
    Gbase.add_nodes_from([0, 1, 2, 3, 4, 5, 6, 7, 8])
    Gbase.add_edges_from(
        [(0, 1), (0, 2), (0, 3), (1, 2), (1, 7), (2, 5), (3, 8), (3, 4), (4, 5), (4, 6), (5, 0),
         (6, 7), (6, 8), (7, 8), (8, 0), (8, 6), (3, 2), (5, 3)])
    i = 0
    for u, v, d in Gbase.edges(data=True):
        d['slots'] = np.zeros(SLOT_NUM)
        d['code'] = i
        i += 1
    return Gbase

def create_gbn_graph():
    Gbase = nx.Graph()
    Gbase.add_nodes_from([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16])
    Gbase.add_edges_from(
        [(0, 2), (0, 8), (1, 2), (1, 3), (1, 4), (2, 4), (3, 4), (3, 9), (4, 8), (4, 10), (4, 9),
         (5, 6), (5, 8), (6, 7), (7, 8), (7, 10), (9, 10), (9, 12), (10, 11), (10, 12), (11, 13),
         (12, 14), (12, 16), (13, 14), (14, 15), (15, 16)])
    i = 0
    for u, v, d in Gbase.edges(data=True):
        d['slots'] = np.zeros(SLOT_NUM)
        d['code'] = i
        i += 1
    return Gbase
def load_graph():
    Gbase = nx.Graph()
    Gbase.add_nodes_from([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13])
    edges=np.load('topo.npy').tolist()
    Gbase.add_edges_from(edges)
    i = 0
    for u, v, d in Gbase.edges(data=True):
        d['slots'] = np.zeros(SLOT_NUM)
        d['code'] = i
        i += 1
    return Gbase

def create_graph(graph_num):
    """
    Generate graphs for training with the same topology.
    """
    if graph_num == 0:
        G = create_nsfnet_graph()
    elif graph_num == 1:
        G = create_geant2_graph()
    elif graph_num == 2:
        G = create_small_top()
    else:
        G = create_gbn_graph()
    return G
def random_destroy_graph():
    Gbase = nx.Graph()
    Gbase.add_nodes_from([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13])
    edges=[(0, 1), (0, 2), (0, 3), (1, 2), (1, 7), (2, 5), (3, 8), (3, 4), (4, 5), (4, 6), (5, 12), (5, 13),
         (6, 7), (7, 10), (8, 9), (8, 11), (9, 10), (9, 12), (10, 11), (10, 13), (11, 12)]
    edges.pop(randrange(len(edges)))
    edges.pop(randrange(len(edges)))
    edges.pop(randrange(len(edges)))
    Gbase.add_edges_from(edges)
    i = 0
    for u, v, d in Gbase.edges(data=True):
        d['slots'] = np.zeros(SLOT_NUM)
        d['code'] = i
        i += 1
    return Gbase
def random_destroy_graph_conn():
    while True:
        g= random_destroy_graph()
        if nx.is_connected(g):
            break
    return g
def random_create_graph():
    while True:
        #随机生成网络  节点数 ，连边概率
        g = nx.gnp_random_graph(20, 0.25)
        if nx.is_connected(g):
            break

    i = 0
    for u, v, d in g.edges(data=True):
        d['slots'] = np.zeros(SLOT_NUM)
        d['code'] = i
        i += 1
    topo = np.array(g.edges)
    np.save('topo.npy', topo)  # 保存为.npy格式
    return g
# 任务用一个三元组表示[head,tail,slot_num]
def create_tasks(node_size):
    res = np.zeros((TASK_NUM+2, 3), dtype=int)
    task_slots = np.random.poisson(TASK_SLOT, TASK_NUM+2)
    #print('业务总带宽='+str( np.sum(task_slots) ))
    for i in range(TASK_NUM+2):
        randoms = np.random.choice(np.array([p for p in range(node_size)]), 2, replace=False)
        head = randoms[0]
        tail = randoms[1]
        task_slot = task_slots[i]
        res[i][0] = head
        res[i][1] = tail
        res[i][2] = task_slot
    np.save('taskscur.npy',res)
    return res
def create_tasks_uneven(node_size):
    res = np.zeros((TASK_NUM, 3), dtype=int)
    task_slots = np.random.poisson(TASK_SLOT, TASK_NUM)
    for i in range(TASK_NUM):
        randoms = np.random.choice(np.array([p for p in range(node_size)]), 2, replace=False,p=[0.2,0.2,0.2,0.1,0.1,0,0,0,0.05,0.05,0.05,0.05,0,0])
        head = randoms[0]
        tail = randoms[1]
        task_slot = task_slots[i]
        res[i][0] = head
        res[i][1] = tail
        res[i][2] = task_slot
    np.save('tasksuneven.npy',res)
    return res

# 查找网络中任意两点间的路由，并按照路由长度递增排序
# 返回一个dict其中 key=head:tail, value=[ [1,2,3],[1,2,5,4,3] ]
def ksp(graph):
    path_record = dict()
    for n1 in graph:
        for n2 in graph:
            if n1 != n2:
                key = str(n1) + ':' + str(n2)
                if key not in path_record:
                    path_record[key] = []
                # cutoff 表示路由长度限制，超过此长度的路由不返回
                [path_record[key].append(p) for p in nx.all_simple_paths(graph, source=n1, target=n2, cutoff=14)]
                path_record[key] = sorted(path_record[key], key=lambda item: (len(item), item))
    for key,value in path_record.items():
        while len(value)<K:
            value.append(value[0])
    return path_record


# 输入路由节点list，输出路由链路编号
def get_link_code(graph, route):
    res = []
    for i in range(len(route) - 1):
        code = graph[int(route[i])][int(route[i + 1])]['code']
        res.append(code)
    return res


# 获取下一状态，以及是否分配成功  如果没有分配成功返回原状态和False
def get_next_state(curr_state1, link_code_list, task_slot):
    curr_state = copy.deepcopy(curr_state1)
    slot_available = util.get_route_slot_final(curr_state, link_code_list)
    # 当前路由缺少网络资源
    if np.max(slot_available) < task_slot:
        curr_state[0][0]=-1
        return curr_state, False
    else:
        index = np.inf
        for idx in range(len(slot_available)):
            if slot_available[idx] >= task_slot:
                index = idx
                break
        for i in link_code_list:
            for j in range(task_slot):
                curr_state[i][index + j] = 1
        return curr_state, True


# 计算边邻接关系
def cal_link_adj(graph):
    # list1 = []
    # list2 = []
    # for u1, v1, d1 in (graph.edges(data=True)):
    #     for u2, v2, d2 in (graph.edges(data=True)):
    #         if (u1 == u2) | (u1 == v2) | (v1 == u2) | (v1 == v2):
    #             list1.append(int(d1['code']))
    #             list2.append(int(d2['code']))
    # res = np.stack((list1, list2), axis=0)
    # return res
    num=graph.number_of_edges()
    res=np.zeros((num,num))
    for u1, v1, d1 in (graph.edges(data=True)):
        for u2, v2, d2 in (graph.edges(data=True)):
            if (u1 == u2) | (u1 == v2) | (v1 == u2) | (v1 == v2):
                res[int(d1['code'])][int(d2['code'])]=1
                res[int(d2['code'])][int(d1['code'])] = 1
    return res


class ENV():
    def __init__(self,task_num,graph_num):
        self.task_num=task_num
        self.graph_num=graph_num
        #self.graph=random_destroy_graph_conn()
        #self.graph = load_graph()
        #self.graph=random_create_graph()
        self.graph=create_graph(graph_num)
        #self.task = np.load('tasks'+'uneven'+'.npy')
        #self.task=create_tasks(self.graph.number_of_nodes())
        #self.task = np.load('tasks' + str(task_num)+ '.npy')
        #self.task=create_tasks_uneven(14)
        self.path_record = ksp(self.graph)
        # 边的邻接关系，为了后续transform
        self.link_adj = cal_link_adj(self.graph)
        self.n_features = SLOT_NUM
        self.n_actions = K
        # 记录当前执行的task编号
        self.task_execute = 0
    def create_task(self):
        self.task = create_tasks(self.graph.number_of_nodes())
    def load_task(self):
        self.task = np.load('tasks' + 'cur' + '.npy')
    def reset(self):
        self.task_execute = 0
        for u, v, d in self.graph.edges(data=True):
            d['slots'] = np.zeros(SLOT_NUM)
    def observe_next_next(self,state,task):
        state=copy.deepcopy(state)
        state_next = np.zeros((K, len(self.graph.edges), SLOT_NUM))
        alloc_flag_next = np.zeros(K)
        for k in range(K):
            route=self.path_record[str(task[0])+':'+str(task[1])][k]
            link_code_list = get_link_code(self.graph, route)
            next_state, alloc_flag = get_next_state(state, link_code_list, task[2])
            state_next[k] = copy.deepcopy(next_state)
            alloc_flag_next[k] = alloc_flag
        return state_next,alloc_flag_next
    # 观察到 当前时刻状态，下一时刻可能的K个状态，边邻接关系矩阵
    # link_msg = 当前网络状态|当前执行任务  link_msg_next=下一网络状态|下一执行任务

    def observe(self):
        link_msg = np.zeros((len(self.graph.edges), SLOT_NUM))
        for u, v, d in self.graph.edges(data=True):
            link_msg[int(d['code'])] = d['slots']
        link_msg_next = np.zeros((K, len(self.graph.edges), SLOT_NUM))
        task = self.task[self.task_execute]
        for k in range(K):
            route = self.path_record[str(task[0]) + ':' + str(task[1])][k]
            link_code_list = get_link_code(self.graph, route)
            next_state, alloc_flag = get_next_state(link_msg, link_code_list, task[2])
            link_msg_next[k] = copy.deepcopy(next_state)
        task = self.task[self.task_execute]
        task_next=self.task[self.task_execute+1]
        res=np.zeros((K+1, len(self.graph.edges), SLOT_NUM+14))
        task=onehot_task(task)
        task_next=onehot_task(task_next)
        cur_state=np.concatenate((link_msg,task),axis=1)
        for i in range(K):
            res[i]= np.concatenate((link_msg_next[i],task_next),axis=1 )
        res[K]=np.concatenate((link_msg,task_next) ,axis=1)
        if res[K][0][0]==-1:
            print('++++++')
        return cur_state,res

    def step(self, state_next,action):
        width = self.task[self.task_execute][2]
        if action==K:
            reward=-width/480
            #reward=0
        else:
            reward = width/480
            for u, v, d in self.graph.edges(data=True):
                d['slots'] = state_next[action][int(d['code'])][:SLOT_NUM]
        self.task_execute+=1
        if self.task_execute>=TASK_NUM:
            done=True
        else:
            done=False
        cur,state_next_next=self.observe()
        return reward,done,state_next_next



if __name__ == "__main__":
    create_tasks_uneven(14)

